# -*- coding: UTF-8 -*- 
# author: notedit
# date: 20110710

import sys
import struct
import bson
import traceback

from twisted.internet import reactor
from twisted.internet import protocol
from twisted.python import log
from twisted.internet import threads
from twisted.application import service,internet

"""mongo protocol"""

QUERY = '$query'
RESULT = 'result'
ADMIN_CMD = 'admin.$cmd'
CMD_MAPPING = {
        'ismaster':{'ismaster':1,'msg':'not paired','ok':1.0},
        }

class MGError(Exception):

    def __init__(self,message,detail):
        self.message = message
        self.detail = detail

    def __str__(self):
        return 'MGError(%s,%s)'%(self.message,self.detail)

    def __repr(self):
        return 'MGError(%s,%s)'%(self.message,self.detail)

class MongoProtocol(protocol.Protocol):

    def __init__(self):
        self.__id = 0
        self.__buffer = ''
        self.__waiting_header = True
        self.__data_length = None
        self.__request_id = 0

    def connectionMade(self):
        print 'connectionMade'

    def connectionLost(self,reason):
        print 'connectionLost'

    def dataReceived(self,data):
        while self.__waiting_header:
            self.__buffer += data
            if len(self.__buffer) < 16:
                break
            header,extra = self.__buffer[:16],self.__buffer[16:]
            self.__buffer = ''
            self.__waiting_header = False
            data_length,request_id,response_to,op_code = struct.unpack('<iiii',header)
            if op_code != 2004:
                print 'OperationError: oocrpc can not support this operation'
                self.transport.loseConnection()
            self.__data_length = data_length - 16
            self.__request_id = request_id
            if extra:
                self.dataReceived(extra)
            break
        else:
            if self.__data_length is not None:
                data,extra = data[:self.__data_length],data[self.__data_length:]
                self.__data_length -= len(data)
            else:
                extra = ''
            self.__buffer += data
            if self.__data_length == 0:
                packet = self.__buffer[:]
                self.__data_length = None
                self.__waiting_header = True
                self.__buffer = ''
                self.processMessage(self.__request_id,packet)
                if extra:
                    self.dataReceived(extra)

    def processMessage(self,request_id,packet):
        try:
            #flags,packet = packet[:4],packet[4:]
            #todo import  bson model
            self.__collection,packet = bson._get_c_string(packet[4:])
            #skip_limit,packet = packet[:8],packet[8:]
            #skip,limit = struct.unpack('<ii',skip_limit)
            spec = bson.BSON(packet[8:]).decode()
            if self.__collection == ADMIN_CMD and spec.get('ismaster',''):
                self.sendMessage(0,bson.BSON.encode(CMD_MAPPING['ismaster']))
            else:
                querystring = spec[QUERY][RESULT]
                func = self.funcmapping.get('query')
                retdict = func(querystring)
                bsonstr = bson.BSON.encode(retdict)
                self.sendMessage(0,bsonstr)
        except MGError,ex:
            traceback.print_exc()
            retdict = {'$err':str(ex)}
            bsonstr = bson.BSON.encode(retdict)
            self.sendMessage(2,bsonstr)
        except Exception,ex:
            traceback.print_exc()
            retdict = {'$err':str(ex)}
            bsonstr = bson.BSON.encode(retdict)
            self.sendMessage(2,bsonstr)

    def sendMessage(self,response_flag,retstr):
        """
        response_flag
            {
                0:CursorNotFound
                1:QueryFailure
            }
        """
        message = struct.pack('<iqii',response_flag,0,0,1) + retstr
        header = struct.pack('<iiii',16 + len(message),0,self.__request_id,1)
        self.transport.write(header + message)

class OocFactory(protocol.ServerFactory):

    protocol = MongoProtocol

    def __init__(self,funcmapping):
        self.funcmapping = funcmapping
        
    def buildProtocol(self,addr):
        p = protocol.ServerFactory.buildProtocol(self,addr)
        p.funcmapping = self.funcmapping
        return p

class OocService(internet.TCPServer):

    def __init__(self,funcmapping,port=27017):
        internet.TCPServer.__init__(self,port,OocFactory(funcmapping))
        return

def start_console(funcmapping,port=27017,threadcount=10):
    log.startLogging(sys.stdout)
    reactor.listenTCP(port,OocFactory(funcmapping))
    reactor.suggestThreadPoolSize(threadcount)
    reactor.run()

def start_daemon(funcmapping,port=27017,threadcount=10):
    app = service.Application('Oocrpc')
    oocservice = OocService(funcmapping,port)
    oocservice.setServiceParent(app)
    reactor.suggestThreadPoolSize(threadcount)
    return app

